'''
研究适配dmc2gym todo
'''

import gymnasium as gym
from PIL import Image
import os

# Initialize the environment
env = gym.make("dmc:Pendulum-swingup-v1")
#

# Reset the environment to get the initial state
state, _ = env.reset()
total_reward = 0
# Run a loop to play the game
for step in range(10000):
    # env.render()  # Render the environment

    # Take a random action
    action = env.action_space.sample()

    # Get the next state, reward, done flag, and info from the environment
    state, reward, done, trunc, info = env.step(action)
    total_reward += reward

    # Save the state as an image
    img = Image.fromarray(state.astype('uint8'), 'RGB')
    img.save(f'state_{step}.png')
    # if reward != 0:
    print("reward: ", reward)
    print("info: ", info)

    # If done, reset the environment
    if done or trunc:
        break

print("Total reward: ", total_reward)

# Close the environment
env.close()